#include "sdf.h"
#include <stdio.h>
#include <string.h>
#include <stdlib.h>

#define TEST_KEY_BITS 256 // 测试 256 位 (32 字节) 会话密钥，GMSSL 默认使用 256

/**
 * 打印字节数组，方便调试
 */
void print_bytes(const char *label, const BYTE *data, ULONG len) {
    printf("%s [%lu bytes]: ", label, len);
    for (ULONG i = 0; i < len; i++) {
        printf("%02X", data[i]);
    }
    printf("\n");
}

// 外部声明：从 sdf.c 访问内部存储
// 注意：这些声明必须与 sdf.h 中的 extern 声明一致
extern uint8_t session_keys[16][32]; 
extern int get_key_index_from_handle(HANDLE hKey);

int main() {
    HANDLE hDevice = NULL;
    LONG ret = SDR_OK;

    // --- 1. 设备打开 ---
    printf("--- 1. Testing SDF_OpenDevice ---\n");
    ret = SDF_OpenDevice(&hDevice);
    if (ret != SDR_OK) {
        printf("Error: SDF_OpenDevice failed. Ret=0x%lX\n", ret);
        return 1;
    }
    printf("SDF_OpenDevice success. Handle=0x%lX\n\n", (ULONG_PTR)hDevice);

    // 假设接收方 R 的内部密钥索引为 2 (用于加密/解密)
    ULONG internal_enc_index = 2; 

    // --- 2. 导出接收方 R 的公钥 (用于发送方封装) ---
    ECCrefPublicKey R_pub_key;
    printf("--- 2. Exporting Encrypt Public Key (Index %lu) ---\n", internal_enc_index);
    ret = SDF_ExportEncPublicKey_ECC(hDevice, internal_enc_index, &R_pub_key);
    if (ret != SDR_OK) {
        printf("Error: SDF_ExportEncPublicKey_ECC failed. Ret=0x%lX\n", ret);
        SDF_CloseDevice(hDevice);
        return 1;
    }
    printf("R Public Key Exported (Bits: %lu)\n", R_pub_key.bits);
    print_bytes("R_pub_key.x", R_pub_key.x, R_pub_key.bits / 8);
    print_bytes("R_pub_key.y", R_pub_key.y, R_pub_key.bits / 8);
    printf("\n");

    // --- 3. 发送方 S: 生成会话密钥并用 R 的公钥加密 (封装) ---
    HANDLE hKey_S = NULL; // 发送方S保存会话密钥的句柄
    ECCCipher key_cipher_text;
    
    printf("--- 3. Generating Session Key with External Public Key (SDF_GenerateKeyWithEPK_ECC) ---\n");
    ret = SDF_GenerateKeyWithEPK_ECC(
        hDevice, 
        TEST_KEY_BITS, 
        0, // uiAlgID, mock中未使用
        &R_pub_key, 
        &key_cipher_text, 
        &hKey_S);

    if (ret != SDR_OK) {
        printf("Error: SDF_GenerateKeyWithEPK_ECC failed. Ret=0x%lX\n", ret);
        SDF_CloseDevice(hDevice);
        return 1;
    }
    printf("Key Encapsulation Success. Sender's Key Handle: 0x%lX\n", (ULONG_PTR)hKey_S);
    print_bytes("Cipher C1 (65B)", key_cipher_text.C1, sizeof(key_cipher_text.C1));
    print_bytes("Cipher C3 (32B)", key_cipher_text.C3, sizeof(key_cipher_text.C3));
    printf("Cipher C2Len: %lu\n", key_cipher_text.C2Len);
    print_bytes("Cipher C2 (Key)", key_cipher_text.C2, key_cipher_text.C2Len);
    printf("\n");

    // --- 4. 接收方 R: 导入会话密钥并用内部私钥解密 (解封装) ---
    HANDLE hKey_R = NULL; // 接收方R导入解密后会话密钥的句柄
    
    printf("--- 4. Importing Session Key with Internal Private Key (SDF_ImportKeyWithISK_ECC) ---\n");
    ret = SDF_ImportKeyWithISK_ECC(
        hDevice, 
        internal_enc_index, // 使用内部私钥索引 2
        &key_cipher_text, 
        &hKey_R);
        
    if (ret != SDR_OK) {
        printf("Error: SDF_ImportKeyWithISK_ECC failed. Ret=0x%lX\n", ret);
        SDF_DestroyKey(hDevice, hKey_S);
        SDF_CloseDevice(hDevice);
        return 1;
    }
    printf("Key Decapsulation Success. Receiver's Key Handle: 0x%lX\n\n", (ULONG_PTR)hKey_R);

    // --- 5. 验证：比较两个句柄指向的密钥数据是否一致 ---
    int index_S = get_key_index_from_handle(hKey_S);
    int index_R = get_key_index_from_handle(hKey_R);
    
    ULONG key_len_bytes = TEST_KEY_BITS / 8;

    if (index_S != -1 && index_R != -1) {
        if (memcmp(session_keys[index_S], session_keys[index_R], key_len_bytes) == 0) {
            printf("--- 5. Verification: SUCCESS ---\n");
            print_bytes("Generated Key", session_keys[index_S], key_len_bytes);
            print_bytes("Decrypted Key", session_keys[index_R], key_len_bytes);
            printf("RESULT: The session keys generated by S and decrypted by R MATCH.\n\n");
        } else {
            printf("--- 5. Verification: FAILED ---\n");
            print_bytes("Generated Key", session_keys[index_S], key_len_bytes);
            print_bytes("Decrypted Key", session_keys[index_R], key_len_bytes);
            printf("RESULT: The session keys DO NOT MATCH.\n\n");
            ret = 1;
        }
    } else {
        printf("Verification FAILED: Could not retrieve internal key data.\n");
        ret = 1;
    }

    // --- 6. 销毁密钥和设备 ---
    printf("--- 6. Destroying Keys and Closing Device ---\n");
    SDF_DestroyKey(hDevice, hKey_S);
    SDF_DestroyKey(hDevice, hKey_R);
    SDF_CloseDevice(hDevice);
    printf("Cleanup complete.\n");

    return ret;
}
